import numpy as np
import cv2
import json
import base64
import os
# labels = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 
#           'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 
#           'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 
#           'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 
#           'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 
#           'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
#           'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
#           'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',
#           'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 
#           'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 
#           'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 
#           'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

labels = ['coal']


def scale_seg_points(label_text,img):
    label_points = {}
    with open(label_text) as f:
        temp_read = f.read().splitlines()
    ii = -1
    for temp_line in temp_read:
        ii += 1
        temp_line = temp_line.split()
        points = np.array(temp_line[1:]).astype(np.float32).reshape(-1,2)
        points[:,0] *= img.shape[1]
        points[:,1] *= img.shape[0]
        # print("points:",points)
        label_points[labels[int(temp_line[0])] + f"_{ii}"] = points
        

        
    return label_points


def yolo_to_labelme_json(save_path,img_path,label_points,hw):
    json_file = {}
    json_file["version"] = "5.4.1"
    json_file["flags"] = {}
    shape_list = []
    for k,v in label_points.items():
        shape_dict = {}
        shape_dict["label"] = k.split("_")[0]
        shape_dict["points"] = v.tolist()
        shape_dict["group_id"] = None
        shape_dict["description"] = ""
        shape_dict["shape_type"] = "polygon"
        shape_dict["flags"] = {}
        shape_dict["mask"] = None
        shape_list.append(shape_dict)
    json_file["shapes"] = shape_list
    json_file["imagePath"] = "abc"
    with open(img_path, 'rb') as image_file:  
        encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
    json_file["imageData"] = encoded_string
    json_file["imageHeight"] = hw[0]
    json_file["imageWidth"] = hw[1]
    # print("json_file：",json_file)

    with open(save_path, 'w') as f:  
        json.dump(json_file, f, indent=4)






if __name__=="__main__":
    img_dir = r"D:\SY_AI_python\yolov5_v7.0\data\predict"
    label_dir = r"D:\SY_AI_python\yolov5_v7.0\runs\predict-seg\exp\labels"


    img_list = os.listdir(img_dir)
    label_list = os.listdir(label_dir)

    for label_file in label_list:
        suffix_name = label_file.rsplit(".",1)[0]
        # print("suffix_name：",suffix_name)
        save_json_path = os.path.join(img_dir,suffix_name + '.json')
        img_path = os.path.join(img_dir,suffix_name + '.jpeg')
        label_path = os.path.join(label_dir,label_file)
        # print("img_path:",img_path)
        img = cv2.imread(img_path)
        hw = img.shape[:2]
        label_points = scale_seg_points(label_path,img) 
        yolo_to_labelme_json(save_json_path,img_path,label_points,hw)
